/* 
 *  Copyright (c) 2020-2021 Xuhpclab. All rights reserved.
 *  Licensed under the MIT License.
 *  See LICENSE file for more information.
 */

#include <map>
#include <string>
#include <sys/stat.h>
#include <math.h>

#include "dr_api.h"
#include "drmgr.h"
#include "drcctlib.h"

using namespace std;

#define DRCCTLIB_PRINTF(_FORMAT, _ARGS...) \
    DRCCTLIB_PRINTF_TEMPLATE("thread_correlation_detection", _FORMAT, ##_ARGS)
#define DRCCTLIB_EXIT_PROCESS(_FORMAT, _ARGS...) \
    DRCCTLIB_CLIENT_EXIT_PROCESS_TEMPLATE("thread_correlation_detection", _FORMAT, ##_ARGS)

#define MAX_CLIENT_CCT_PRINT_DEPTH 10

static std::string g_folder_name;
static int tls_idx;

typedef struct _cache_contention_t {
    int miss_count;
    int inval_count;
} cache_contention_t;

typedef struct _per_thread_t {
    int pt_id;
    long long int thread_local_miss_count;
    long long int thread_local_invalidation_count;
    long long int thread_local_cold_miss_count;
    file_t output_file;
// #define DEBUG_CACHE_CONTENTION
#ifdef DEBUG_CACHE_CONTENTION
    file_t log_file;
#endif
} per_thread_t;

//map<int, cache_contention_t> *thread_count_map = new map<int, cache_contention_t>();
static int total_miss_count = 0;
static int total_invalidation_count = 0;
static int total_cold_miss_count = 0;

typedef struct Thread_array {
    int array[32];
} T_array;

map<int, T_array> *thread_correlation_count_map = new map<int, T_array>;

typedef struct _top_N_format {
    int i;
    int j;
    int times;
} top_N_format;

//定义所要模拟的cache line大小，单位是字节（byte�?
#define CacheLineSize   64

//定义将访存地址转换为cache line对齐�?
//（是否可以考虑修改drcctlib.cpp中的fwt_init_shadow_memory函数？只需要每个开始地址存储一个unit即可，因为保证访存地址进行cache line align�?
//目前是以cachelinesize=64计算的align，最�?位清�?
#define CacheLineAlign(_addr)   _addr&0xffffffffffffffc0

//return true means happens cache miss
//return false means goes on, nothing happens
inline bool
test_miss(cache_cont_t cur_bitmap, int pt_id)
{
    return !((cur_bitmap.bitmap & (1 << pt_id)) > 0);
}

//return true means happens cache invalidation
//return false means goes on, nothing happens
inline bool
test_inval(cache_cont_t cur_bitmap, int pt_id)
{
    return (cur_bitmap.bitmap & (~(1 << pt_id))) > 0;
}

//use OR instruction to add a thread into bitmap
inline void
after_miss_happens_refresh_bitmap(cache_cont_t* cur_bitmap, int pt_id)
{
    (*cur_bitmap).bitmap = (*cur_bitmap).bitmap | (1 << pt_id);
    return ;
}

//use MOV instruction to clear all other thread flag, and set its exclusive own flag
inline void
after_invalidation_happens_refresh_bitmap(cache_cont_t* cur_bitmap, int pt_id)
{
    (*cur_bitmap).bitmap = 1 << pt_id;
    return ;
}

/*
inline void
update_miss_count(int pt_id)
{
    map<int, cache_contention_t>::iterator it = (*thread_count_map).find(pt_id);
    if (it != (*thread_count_map).end()) {
        it->second.miss_count++;
    } else {
        cache_contention_t new_cont;
        new_cont.miss_count = 1;
        new_cont.inval_count = 0;
        (*thread_count_map).insert(pair<int, cache_contention_t>(pt_id, new_cont));
    }
    total_miss_count++;
    return ;
}

inline void
update_inval_count(int pt_id) 
{
    map<int, cache_contention_t>::iterator it = (*thread_count_map).find(pt_id);
    if (it != (*thread_count_map).end()) {
        it->second.inval_count++;
    } else {
        cache_contention_t new_cont;
        new_cont.miss_count = 0;
        new_cont.inval_count = 1;
        (*thread_count_map).insert(pair<int, cache_contention_t>(pt_id, new_cont));
    }
    total_invalidation_count++;
    return ;
}
*/

inline void 
update_RLUTI_bitmap_after_store(cache_cont_t* cur_bitmap, int pt_id)
{
    (*cur_bitmap).RLUTI_bitmap = 1 << pt_id;
    return ;
}

inline void
update_thread_correlation_info(cache_cont_t cur_bitmap, int pt_id)
{
    int RLUTI_bitmap = cur_bitmap.RLUTI_bitmap;
    int another_pt_id = 0;
    //get another_pt_id
    for (another_pt_id = 0; ;another_pt_id++) {
        if (RLUTI_bitmap > 1) { //1 = 2^0
            RLUTI_bitmap = RLUTI_bitmap >> 1;
        } else {
            break;
        }
    }
    //dr_printf("another pt id is %d\n\n", another_pt_id);
    map<int, T_array>::iterator it1 = (*thread_correlation_count_map).find(pt_id);
    map<int, T_array>::iterator it2 = (*thread_correlation_count_map).find(another_pt_id);
    if (it1 != (*thread_correlation_count_map).end()) {
        it1->second.array[another_pt_id]++;
    } else {
        T_array new_array;
        for (int i = 0; i < 32; i++) {
            new_array.array[i] = 0;
        }
        (*thread_correlation_count_map).insert(pair<int, T_array>(pt_id, new_array));
    }

    if (it2 != (*thread_correlation_count_map).end()){
        it2->second.array[pt_id]++;
    } else {
        T_array new_array;
        for (int i = 0; i < 32; i++) {
            new_array.array[i] = 0;
        }
        (*thread_correlation_count_map).insert(pair<int, T_array>(another_pt_id, new_array));
    }
    return ;
}


// client want to do
inline void
DoWhatClientWantTodo(void *drcontext, per_thread_t *pt, context_handle_t cur_ctxt_hndl,
                     app_pc cur_addr, char type, ushort size)
{
    //获取当前运行线程的id
   int cur_thread_id = drcctlib_get_thread_id();
   //地址对齐
   uint64_t aligned_addr = CacheLineAlign((uint64_t) cur_addr);
   cache_cont_t* cur_bitmap;

   //检测访存空间的大小，若访存数据大小超过一个Cache Line的大小，则需要分多个Cache Line处理
   //Cache_Line_nums代表该次访存涉及几个Cache Line
   int Cache_Line_nums = 0;
   if (((uint64_t)cur_addr + size) <= (aligned_addr + 64)) {
       Cache_Line_nums = 1;
   } else {
       int size_remain = size - (aligned_addr + 64 - (uint64_t)cur_addr);
       Cache_Line_nums = (int)ceil((double)size_remain / (double)CacheLineSize);
   }

   //dr_printf("3: in DoWhatClientWantTodo, before for\n");
   
   for (int i = 0; i < Cache_Line_nums; i++) {
        //初次访问探测机制
        if(!fwt_test_shadow_memory((size_t)aligned_addr)) {
            //未初始化过，需进行init
            pt->thread_local_cold_miss_count++;
            //total_cold_miss_count++;
            fwt_init_shadow_memory((void*)aligned_addr, CacheLineSize);
            cur_bitmap = fwt_get_shadow_memory((void*) aligned_addr);
	          after_miss_happens_refresh_bitmap(cur_bitmap, cur_thread_id);
        }

        //对于已经有内容的shadow memory地址，将其存储的信息取出
        cur_bitmap = fwt_get_shadow_memory((void*) aligned_addr);
        
        //dr_printf("access addr is : %x, cur_thread_id is %d \n", aligned_addr, cur_thread_id);

        if(type == 0) {
            //说明该内存访问操作是读操作，检测cache miss
            if (test_miss(*cur_bitmap, cur_thread_id)) {
                pt->thread_local_miss_count++;
                //update_miss_count(cur_thread_id);
                after_miss_happens_refresh_bitmap(cur_bitmap, cur_thread_id);
                //发生cache miss后，检测当前的线程交互情况
                update_thread_correlation_info(*cur_bitmap, cur_thread_id);
            }
        } else {
        //dr_printf("write addr is : %x, cur_thread_id is %d \n", aligned_addr, cur_thread_id);
            //说明该内存访问操作是写操作，检测cache invalidation
            if (test_inval(*cur_bitmap, cur_thread_id)) {
            //dr_printf("happens inval! cur_bitmap is : %d, cur_thread_id is %d\n", cur_bitmap->bitmap, cur_thread_id);
                pt->thread_local_invalidation_count++;
                //update_inval_count(cur_thread_id);
                after_invalidation_happens_refresh_bitmap(cur_bitmap, cur_thread_id);
                //发生cache invalidation后，检测当前的线程交互情况
                update_thread_correlation_info(*cur_bitmap, cur_thread_id);
            }
            //对于store操作，更新Cache Line的RLUTI bitmap
            update_RLUTI_bitmap_after_store(cur_bitmap, cur_thread_id);
        }

        //对于一次访存空间大于一个Cache Line Size的情况，循环判断全部地址
        //当前Cache Line判断结束，地址增加512bit，即64 bytes，即为下一个Cache Line的起始地址
        aligned_addr += 64;
   }
}

static inline void
InstrumentPerBBCache(void *drcontext, context_handle_t ctxt_hndl, int32_t slot_num,
                     int32_t mem_ref_num, mem_ref_msg_t *mem_ref_start, void **data)
{
    per_thread_t *pt;
    if (*data != NULL) {
        pt = (per_thread_t *)*data;
    } else {
        pt = (per_thread_t *)drmgr_get_tls_field(drcontext, tls_idx);
        *data = pt;
    }

    for (int32_t i = 0; i < mem_ref_num; i++) {
        if (mem_ref_start[i].slot >= slot_num) {
            break;
        }
        if (read_main_flag() == 0)
            continue;
        //dr_printf("2: in InstrumentPerBBCache, before DoWhatClientWantTodo\n");
        DoWhatClientWantTodo(drcontext, pt, ctxt_hndl + mem_ref_start[i].slot,
                             mem_ref_start[i].addr, mem_ref_start[i].type, mem_ref_start[i].size);
    }
}

#ifdef DEBUG_CACHE_CONTENTION
static void
ThreadDebugFileInit(per_thread_t *pt)
{
    int32_t id = drcctlib_get_thread_id();
    char name[MAXIMUM_FILEPATH] = "";
    sprintf(name + strlen(name), "%s/thread-%d.debug.log", g_folder_name.c_str(), id);
    pt->log_file = dr_open_file(name, DR_FILE_WRITE_APPEND | DR_FILE_ALLOW_LARGE);
    DR_ASSERT(pt->log_file != INVALID_FILE);
}
#endif

static void
ThreadOutputFileInit(per_thread_t *pt)
{
    int32_t id = drcctlib_get_thread_id();
    char name[MAXIMUM_FILEPATH] = "";
    sprintf(name + strlen(name), "%s/thread-%d.topn.log", g_folder_name.c_str(), id);
    pt->output_file = dr_open_file(name, DR_FILE_WRITE_OVERWRITE | DR_FILE_ALLOW_LARGE);
    DR_ASSERT(pt->output_file != INVALID_FILE);
}

static void
ClientThreadStart(void *drcontext)
{
    per_thread_t *pt = (per_thread_t *)dr_thread_alloc(drcontext, sizeof(per_thread_t));
    if (pt == NULL) {
        DRCCTLIB_EXIT_PROCESS("pt == NULL");
    }
    drmgr_set_tls_field(drcontext, tls_idx, (void *)pt);
    pt->pt_id = drcctlib_get_thread_id();
    pt->thread_local_miss_count = 0;
    pt->thread_local_invalidation_count = 0;
    pt->thread_local_cold_miss_count = 0;
    ThreadOutputFileInit(pt);
#ifdef DEBUG_CACHE_CONTENTION
    ThreadDebugFileInit(pt);
#endif
}

void
PrintTopN(per_thread_t *pt, int N) 
{
    top_N_format top[N];
    for (int i = 0; i < N; i++) {
        top[i].i = 0;
        top[i].j = 0;
        top[i].times = 0;
    }
    for (int i = 0; i < 32; i++) {
        map<int, T_array>::iterator it = (*thread_correlation_count_map).find(i);
        if (it != (*thread_correlation_count_map).end()) {
            for (int j = i + 1; j < 32; j++) { //�����߳�i�뱾�����޽�����ϵ���������j = i + 1
                if (it->second.array[j] <= top[N - 1].times) { //�����ǰԪ�ر�top��������С�Ļ�С���������κβ���
                    continue;
                }
                for (int k = 0; k < N; k++) {
                    if (it->second.array[j] > top[k].times) {
                        for (int p = N - 1; p > k; p--) {
                            top[p].times = top[p - 1].times;
                            top[p].i = top[p - 1].i;
                            top[p].j = top[p - 1].j;
                        }
                        top[k].times = it->second.array[j];
                        top[k].i = i;
                        top[k].j = j;
                        break;
                    }
                }
            }
        }
    }
    for (int i = 0; i < N; i++) {
      dr_fprintf(pt->output_file, "No %d. : thread <%d, %d> correlates %d times.\n", i, top[i].i, top[i].j, top[i].times);
    }
}

void
PrintResult(per_thread_t *pt) 
{
    cache_contention_t cur_cont;
    int cur_thread_id = pt->pt_id;//drcctlib_get_thread_id();
    dr_fprintf(pt->output_file, "current thread id is %d\n ===========================================\n\n", cur_thread_id);
    /*map<int, cache_contention_t>::iterator it = (*thread_count_map).find(cur_thread_id);
    if (it == (*thread_count_map).end()) {
        dr_fprintf(pt->output_file, "current thread doesn't exist! Something goes wrong!!\n");
        return ;
    }*/
    dr_fprintf(pt->output_file, "thread miss count : %d  total miss count : %d  percentage : %2f% \n\n", 
                                                                                pt->thread_local_miss_count, total_miss_count, 
                                                                                    ((double)pt->thread_local_miss_count / total_miss_count) * 100);
    dr_fprintf(pt->output_file, "thread inval count : %d  total inval count : %d  percentage : %2f% \n\n", 
                                                                                pt->thread_local_invalidation_count, total_invalidation_count,
                                                                                    ((double)pt->thread_local_invalidation_count / total_invalidation_count) * 100);
    dr_fprintf(pt->output_file, "================================\n thread correlation info \n=====================================\n");
    PrintTopN(pt, 3);
   map<int, T_array>::iterator it1 = (*thread_correlation_count_map).find(cur_thread_id);
   if (it1 == (*thread_correlation_count_map).end()) {
       dr_fprintf(pt->output_file, "current thread doesn't have correlation infomation!!!\n");
       return ;
   }
   for (int i = 0; i < 32; i++) {
       if (i != cur_thread_id) {
            dr_fprintf(pt->output_file, "current thread <%d> has %d times correlation with thread <%d>\n", cur_thread_id, it1->second.array[i], i);
       }
   }
}

static void
ClientThreadEnd(void *drcontext)
{
    per_thread_t *pt = (per_thread_t *)drmgr_get_tls_field(drcontext, tls_idx);
    total_cold_miss_count += pt->thread_local_cold_miss_count;
    total_miss_count += pt->thread_local_miss_count;
    total_invalidation_count += pt->thread_local_invalidation_count;
    PrintResult(pt);
    dr_close_file(pt->output_file);
#ifdef DEBUG_CACHE_CONTENTION
    dr_close_file(pt->log_file);
#endif
    dr_thread_free(drcontext, pt, sizeof(per_thread_t));
}

static void
ClientInit(int argc, const char *argv[])
{
    char name[MAXIMUM_FILEPATH] = "";
    DRCCTLIB_INIT_LOG_FILE_NAME(
        name, "drcctlib_thread_correlation_detection", "out");
    g_folder_name.assign(name, strlen(name));
    mkdir(g_folder_name.c_str(), S_IRWXU | S_IRWXG | S_IROTH | S_IXOTH);
}

static void
ClientExit(void)
{
    drcctlib_exit();
    if (!drmgr_unregister_thread_init_event(ClientThreadStart) ||
        !drmgr_unregister_thread_exit_event(ClientThreadEnd) ||
        !drmgr_unregister_tls_field(tls_idx)) {
        DRCCTLIB_PRINTF(
            "ERROR: drcctlib_thread_correlation_detection failed to unregister in ClientExit");
    }
    drmgr_exit();
}

#ifdef __cplusplus
extern "C" {
#endif

DR_EXPORT void
dr_client_main(client_id_t id, int argc, const char *argv[])
{
    dr_set_client_name("DynamoRIO Client 'drcctlib_thread_correlation_detection'",
                       "http://dynamorio.org/issues");
    ClientInit(argc, argv);

    if (!drmgr_init()) {
        DRCCTLIB_EXIT_PROCESS(
            "ERROR: drcctlib_thread_correlation_detection unable to initialize drmgr");
    }
    drmgr_priority_t thread_init_pri = { sizeof(thread_init_pri),
                                         "drcctlib_thread_correlation-thread_init", NULL, NULL,
                                         DRCCTLIB_THREAD_EVENT_PRI + 1 };
    drmgr_priority_t thread_exit_pri = { sizeof(thread_exit_pri),
                                         "drcctlib_thread_correlation-thread-exit", NULL, NULL,
                                         DRCCTLIB_THREAD_EVENT_PRI + 1 };
    drmgr_register_thread_init_event_ex(ClientThreadStart, &thread_init_pri);
    drmgr_register_thread_exit_event_ex(ClientThreadEnd, &thread_exit_pri);
    tls_idx = drmgr_register_tls_field();
    if (tls_idx == -1) {
        DRCCTLIB_EXIT_PROCESS(
            "ERROR: drcctlib_thread_correlation_detection drmgr_register_tls_field fail");
    }
    //dr_printf("1: before drcctlib_init_ex\n");
    drcctlib_init_ex(DRCCTLIB_FILTER_MEM_ACCESS_INSTR, INVALID_FILE, NULL, NULL,
                     InstrumentPerBBCache,
                     DRCCTLIB_COLLECT_DATA_CENTRIC_MESSAGE | DRCCTLIB_CACHE_MODE |
                         DRCCTLIB_CACHE_MEMEORY_ACCESS_ADDR);
    dr_register_exit_event(ClientExit);
}

#ifdef __cplusplus
}
#endif
